// Copyright (C) 2011  Andrew H. Chan
// 
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 2
// of the License, or (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

#include "sampler.h"
#include "util.h"

using namespace std;

extern gsl_rng *rg;

typedef std::list<Table> table_t;
typedef std::pair<std::vector<int>, table_t::iterator> hap_table_t;
typedef std::vector<hap_table_t> hap_t;

Sampler::Sampler(Input_data const& input_data) : num_haps(input_data.num_haps), segsites(input_data.segsites) {
    data = input_data;
}

int Sampler::gibbs_sample(uint iterations) {
    const uint burn_in = 2000;
    double tau = 0.01 * num_haps/2; // 1% of the population size
    double alpha_h = 10.0;
    double beta_h = .1;
    double alpha_g = 10.0;
    double beta_g = 0.01;

    std::vector<std::vector<int> > const &genotypes = data.genotypes;

    comparison_num_mistakes.resize(num_haps, 0);
    hap_counts0.resize(num_haps);
    hap_counts1.resize(num_haps);
    inf_haps.resize(num_haps);
    for(uint i=0; i < num_haps; ++i) {
        hap_counts0[i].resize(segsites, 0); 
        hap_counts1[i].resize(segsites, 0); 
        inf_haps[i].resize(segsites, 0);
    }
    posterior_counts.resize(num_haps/2);

    // compute gen_het_index
    gen_het_index.resize(num_haps/2);
    gen_het_num.resize(num_haps/2);
    for(uint i=0; i<num_haps/2; ++i) {
        int het_num = 0;
        for(uint j=0; j<segsites; ++j) {
            if(genotypes[i][j] == 2) {
                ++het_num;
                gen_het_index[i].push_back(j);
            }
        }
        gen_het_num[i] = het_num;
    }

    //debug
    size_t max_num_tables = 0;

    // ancestral haplotypes
    table_t tables;
    // haplotypes and which table their sitting at
    hap_t haplotypes;

    // initialize tables
    tables.push_back(Table(vector<int>(segsites, 0)));
    cout << "segsites: " << segsites << endl; 

    // initialize haplotypes
    for(uint i=0; i<num_haps; ++i) {
        haplotypes.push_back(hap_table_t(vector<int>(segsites, 0), tables.begin()));
    }
    // compute number of non mutations
    // since all haplotypes are all 0s, there are no mutations
    for(uint j=0; j<segsites; ++j) {
        tables.front().nnm[j] = num_haps;
    }
    tables.front().num_at_table = num_haps;

    // Gibbs updating
    for(uint iteration=0; iteration < iterations; ++iteration) {
        max_num_tables = max(max_num_tables, tables.size());
        ////////////////////////////////////////////
        // update tables the haplotypes sit at
        ////////////////////////////////////////////
        // use Metropolis-Hastings
        for(hap_t::iterator hap_iter = haplotypes.begin(); hap_iter != haplotypes.end(); ++hap_iter) {
            table_t::iterator cur_table = hap_iter->second;

            // temporarily subtract cur haplotype from his table
            // must ignore the current haplotype when computing proposal
            --cur_table->num_at_table;
            assert(cur_table->num_at_table >= 0);
            for(uint s_index=0; s_index < segsites; ++s_index) {
                if(cur_table->ancestor[s_index] == hap_iter->first[s_index]) {
                    --cur_table->nnm[s_index];
                    assert(cur_table->nnm[s_index] >= 0);
                }
            }

            ///////////////////////////////////////////////////////////
            // propose to sit at a new table
            // go through each table
            double running_sum = 0.0;
            double denom = (double)num_haps - 1.0 + tau;
            //cout << "go through each table" << endl;
            // Get uniform random number between 0 and 1
            double unif = gsl_rng_uniform(rg);
            table_t::iterator table_to_sit_at;
            for(table_to_sit_at=tables.begin(); table_to_sit_at != tables.end(); ++table_to_sit_at) {
                double numtable = table_to_sit_at->num_at_table;
                running_sum += numtable / denom;
                if(unif < running_sum) {
                    break;
                }
            }

            // update mutation counts for each seg site for table he left 
            // compute likelihood of current haplotype from the table
            //    he just left, conditioned on everyone there except for him.
            // This is for the acceptance probability.
            // I don't need -1 here because I've already removed the
            //    current haplotype.
            double accept_bottom = 1.0;
            for(uint j=0; j<segsites; ++j) {
                // compute denominator of accept prob
                if(hap_iter->first[j] == cur_table->ancestor[j]) {
                    accept_bottom *= (alpha_h + cur_table->nnm[j]) / (alpha_h + beta_h + cur_table->num_at_table);
                } else {
                    accept_bottom *= (beta_h + cur_table->num_at_table - cur_table->nnm[j]) / (alpha_h + beta_h + cur_table->num_at_table);
                }
            }
            vector<int> new_ancestor;
            vector<int> new_nnm;
            if(table_to_sit_at == tables.end()) {
                // start a new table
                // sample new table from equation (19a)
                new_ancestor.clear();
                for(uint j=0; j < segsites; ++j) {
                    double unif = gsl_rng_uniform(rg);
                    double thr = alpha_h / (alpha_h + beta_h);
                    if(unif < thr) {
                        new_ancestor.push_back(hap_iter->first[j]);
                        new_nnm.push_back(1);
                    } else {
                        new_ancestor.push_back(1-hap_iter->first[j]);
                        new_nnm.push_back(0);
                    }
                }
            } else {
                // sample existing table from equation(19b)
                // sit at existing table
                // temporarily add haplotype to existing table
                ++table_to_sit_at->num_at_table;
                // recompute nnm for existing table
                for(uint j=0; j<segsites; ++j) {
                    if(hap_iter->first[j] == table_to_sit_at->ancestor[j]) {
                        ++table_to_sit_at->nnm[j];
                    }
                }
                int nmwrt0;
                int mwrt0;
                int nmwrt1;
                int mwrt1;
                new_ancestor.clear();
                for(uint j=0; j < segsites; ++j) {
                    if(table_to_sit_at->ancestor[j] == 0) {
                        mwrt1 = nmwrt0 = table_to_sit_at->nnm[j];
                        nmwrt1 = mwrt0 = table_to_sit_at->num_at_table - nmwrt0;
                        assert(nmwrt0 >= 0);
                        assert(mwrt0 >= 0);
                    } else {
                        assert(table_to_sit_at->ancestor[j] == 1);
                        mwrt0 = nmwrt1 = table_to_sit_at->nnm[j];
                        nmwrt0 = mwrt1 = table_to_sit_at->num_at_table - nmwrt1;
                        assert(nmwrt1 >= 0);
                        assert(mwrt1 >= 0);
                    }
                    // compute gamma functions
                    double numerator = gsl_sf_gamma(alpha_h + (double)nmwrt0) * gsl_sf_gamma(beta_h + (double)mwrt0);
                    double denominator = numerator + gsl_sf_gamma(alpha_h + (double)nmwrt1)*gsl_sf_gamma(beta_h + (double)mwrt1);
                    double probability = numerator / denominator;
                    double unif = gsl_rng_uniform(rg);
                    if(unif < probability) {
                        // anc_j is 0
                        new_ancestor.push_back(0);
                        new_nnm.push_back(nmwrt0);
                    } else {
                        // anc_j is 1
                        new_ancestor.push_back(1);
                        new_nnm.push_back(nmwrt1);
                    }
                }
            }
            assert(new_nnm.size() == segsites);
            assert(new_ancestor.size() == segsites);

            // compute acceptance probability
            // compute numerator
            double accept_top = 1.0;
            if(table_to_sit_at == tables.end()) {
                // if completely new table
                for(uint j=0; j<segsites; ++j) {
                    if(hap_iter->first[j] == new_ancestor[j]) {
                        accept_top *= alpha_h/(alpha_h + beta_h);
                    } else {
                        accept_top *= beta_h/(alpha_h + beta_h);
                    }
                }
            } else {
                // existing table
                for(uint j=0; j<segsites; ++j) {
                    // remember to subtract out the current haplotype
                    if(hap_iter->first[j] == new_ancestor[j]) {
                        accept_top *= (alpha_h + new_nnm[j] - 1) / (alpha_h + beta_h + table_to_sit_at->num_at_table-1);
                    } else {
                        assert(hap_iter->first[j] != new_ancestor[j]);
                        accept_top *= (beta_h + table_to_sit_at->num_at_table - new_nnm[j] - 1) / (alpha_h + beta_h + table_to_sit_at->num_at_table - 1);
                    }
                }
            }
            // compute denominator 
            // denominator was already computed
            unif = gsl_rng_uniform(rg);
            // accept or reject?
            if(unif < accept_top / accept_bottom)  {
                // if accept,
                // if table is new, then add table to tables list
                if(table_to_sit_at == tables.end()) {
                    tables.push_back(Table(new_ancestor));
                    tables.back().nnm = new_nnm;
                    tables.back().num_at_table = 1;
                    hap_iter->second = --tables.end(); 
                } else {
                    // if table is existing, then update the ancestor 
                    table_to_sit_at->ancestor = new_ancestor;
                    table_to_sit_at->nnm = new_nnm;
                    hap_iter->second = table_to_sit_at;
                    // num_at_table was already updated
                }
                // remove old table if necessary (num_at_table == 0)
                if(cur_table->num_at_table == 0) {
                    tables.erase(cur_table);
                }
            } else {
                // if reject, add haplotype back to his old table
                // he's already pointing to his old table
                // increment old table's hap counter
                ++cur_table->num_at_table;
                // update nnm for his old table
                for(uint s_index=0; s_index < segsites; ++s_index) {
                    if(cur_table->ancestor[s_index] == hap_iter->first[s_index]) {
                        ++cur_table->nnm[s_index];
                        assert(cur_table->nnm[s_index] >= 0);
                    }
                }

                // if he proposed to sit at an existing table,
                //     remove him from that table
                if(table_to_sit_at != tables.end()) {
                    // minus 1 from num_at_table
                    --table_to_sit_at->num_at_table;
                    // recompute nnm (remove the haplotype)
                    for(uint s_index=0; s_index < segsites; ++s_index) {
                        if(table_to_sit_at->ancestor[s_index] == hap_iter->first[s_index]) {
                            --table_to_sit_at->nnm[s_index];
                            assert(table_to_sit_at->nnm[s_index] >= 0);
                        }
                    }
                }
                // if he proposed to sit at a new table and rejects,
                //    then nothing needs to be done except to add him
                //    back to his old table (which has already been done).
            } // end reject
        } // end iterating through haplotypes, updating their tables

        // update haplotypes based on genotypes
        // compute u, u', and u''
        // add/subtract out differences later
        int u = 0; // matches with genotype
        int up = 0; // single difference with genotype
        int udp = 0; // double difference with genotype

        // iterate through genotypes
        for(uint i=0; i < num_haps/2; ++i) {
            for(uint j=0; j < segsites; ++j) {
                if(genotypes[i][j] == 0) { // 00
                    if(haplotypes[2*i].first[j] == 0 && haplotypes[2*i+1].first[j] == 0) {
                        ++u; // perfect match: 00 - 00
                    } else if((haplotypes[2*i].first[j] == 0 && haplotypes[2*i+1].first[j] == 1) || (haplotypes[2*i].first[j] == 1 && haplotypes[2*i+1].first[j] == 0) ) {
                        ++up; // single difference 01 - 00 or 10 - 00
                    } else {
                        assert(haplotypes[2*i].first[j] == 1 && haplotypes[2*i+1].first[j] == 1);
                        ++udp; // double difference 11 - 00
                    }
                } else if(genotypes[i][j] == 1) { // 11
                    if(haplotypes[2*i].first[j] == 0 && haplotypes[2*i+1].first[j] == 0) {
                        ++udp; // double difference 11 - 00
                    } else if((haplotypes[2*i].first[j] == 0 && haplotypes[2*i+1].first[j] == 1) || (haplotypes[2*i].first[j] == 1 && haplotypes[2*i+1].first[j] == 0)) {
                        ++up; // single difference 01 - 11 or 10 - 11
                    } else {
                        assert(haplotypes[2*i].first[j] == 1 && haplotypes[2*i+1].first[j] == 1);
                        ++u; // perfect match: 11 - 11
                    }
                } else {
                    assert(genotypes[i][j] == 2); // {0,1} 
                    if(haplotypes[2*i].first[j] == 0 && haplotypes[2*i+1].first[j] == 0) {
                        ++up; // single difference 00 - {0,1}
                    } else if((haplotypes[2*i].first[j] == 0 && haplotypes[2*i+1].first[j] == 1) || (haplotypes[2*i].first[j] == 1 && haplotypes[2*i+1].first[j] == 0)) {
                        ++u; // perfect match 01 - {0,1} or 10 - {0,1}
                    } else {
                        assert(haplotypes[2*i].first[j] == 1 && haplotypes[2*i+1].first[j] == 1);
                        ++up; // single difference 11 - {0,1}
                    }
                }
            }
        }

        int hap_counter = 0;
        // iterate through haplotypes
        for(hap_t::iterator hap_iter = haplotypes.begin(); hap_iter != haplotypes.end(); ++hap_iter) {
            hap_t::iterator other_hap;
            if(hap_counter%2 == 1) { // if odd, then other is 1 below
                other_hap = hap_iter;
                --other_hap;
            } else { // if even, then other is 1 above
                other_hap = hap_iter;
                ++other_hap;
            }

            for(uint j=0; j < segsites; ++j) {
                // get temporary u's
                int tmp_u0 = u; // u if hap_ij becomes 0
                int tmp_up0 = up; // up if hap_ij becomes 0
                int tmp_udp0 = udp; // udp if hap_ij becomes 0
                int tmp_u1 = u; // u if hap_ij becomes 1
                int tmp_up1 = up; // up if hap_ij becomes 1
                int tmp_udp1 = udp; // udp if hap_ij becomes 1
                if(genotypes[hap_counter/2][j] == 0) { // 00
                    if(hap_iter->first[j] == 1) { // old allele is 1
                        if(other_hap->first[j] == 0) { // old is 10
                            // if become 0, thus 00
                            --tmp_up0; // lose single difference
                            ++tmp_u0; // gain perfect match
                            // if become 1, no difference (already 1)
                            // so tmp_ux1 remains the same
                        } else { // old is 11
                            // if become 0, thus 01
                            --tmp_udp0; // lose double difference
                            ++tmp_up0; // gain single difference
                            // if become 1, no difference (already 1)
                            // so tmp_ux1 remains the same
                        }
                    } else { // old allele is 0
                        assert(hap_iter->first[j] == 0);
                        if(other_hap->first[j] == 0) { // old is 00
                            // if become 1, thus 10
                            --tmp_u1; // lose perfect match
                            ++tmp_up1; // gain single difference
                            // if become 0, no difference
                        } else { // old is 01
                            // if become 1, thus 11
                            --tmp_up1; // lose single difference
                            ++tmp_udp1; // gain double difference
                            // if become 0, no difference
                        }
                    }
                } else if(genotypes[hap_counter/2][j] == 1) { // genotype 11
                    if(hap_iter->first[j] == 1) { // old allele is 1
                        if(other_hap->first[j] == 0) { // old is 10
                            // become 0, thus 00
                            --tmp_up0; // lose single difference
                            ++tmp_udp0; // gain double difference
                        } else { // old is 11
                            // become 0, thus 01 
                            --tmp_u0; // lose perfect match
                            ++tmp_up0; // gain single difference
                        }
                    } else { // old allele is 0
                        if(other_hap->first[j] == 0) { // old is 00
                            // become 1, thus 10
                            --tmp_udp1; // lose double difference
                            ++tmp_up1; // gain single difference
                        } else { // old is 01
                            // become 1, thus 11
                            --tmp_up1; // lose single difference
                            ++tmp_u1; // gain perfect match
                        }
                    }
                } else { // genotype {0,1}
                    assert(genotypes[hap_counter/2][j] == 2);
                    if(hap_iter->first[j] == 1) { // old allele is 1
                        if(other_hap->first[j] == 0) { // old is 10
                            // become 0, thus 00
                            --tmp_u0; // lose perfect match
                            ++tmp_up0; // gain single difference
                        } else { // old is 11
                            // become 0, thus 01
                            --tmp_up0; // lose single difference
                            ++tmp_u0; // gain perfect match
                        }
                    } else { // old allele is 0
                        if(other_hap->first[j] == 0) { // old is 00
                            // become 1, thus 10
                            --tmp_up1; // lose single difference
                            ++tmp_u1; // gain perfect match
                        } else { // old is 01 
                            // become 1, thus 11
                            --tmp_u1; // lose perfect match
                            ++tmp_up1; // gain single difference
                        }
                    }
                }

                double left_denom_first = 1.0;
                if(tmp_u0 >= tmp_u1) {
                    for(int kk = tmp_u1; kk < tmp_u0; ++kk) {
                        left_denom_first *= ((double)kk + alpha_g);
                    }
                    left_denom_first = 1.0/left_denom_first;
                }
                else {
                    for(int kk = tmp_u0; kk < tmp_u1; ++kk) {
                        left_denom_first *= ((double)kk + alpha_g);
                    }
                }

                double left_denom_second = 1.0;
                if((tmp_up0 + tmp_udp0) >= (tmp_up1 + tmp_udp1)) {
                    for(int kk = tmp_up1 + tmp_udp1; kk < tmp_up0 + tmp_udp0; ++kk) {
                        left_denom_second *= ((double)kk + beta_g);
                    }
                    left_denom_second = 1.0/left_denom_second;
                }
                else {
                    for(int kk = tmp_up0 + tmp_udp0; kk < tmp_up1 + tmp_udp1; ++kk) {
                        left_denom_second *= ((double)kk + beta_g);
                    }
                }

                double mumu = pow(0.5, (double)tmp_up1 - (double)tmp_up0) * pow(0.5, (double)tmp_udp1 - (double)tmp_udp0);
                double leftside0 = 1.0 / (left_denom_first * left_denom_second * mumu + 1.0);
                assert(leftside0 >= 0.0 && leftside0 <= 1.0);

                double leftside1 = 1.0 - leftside0;

                ////////////////////////////////////////////////////////////////
                int num_non_muts0; // num non mutations if become 0
                int num_non_muts1; // num non mutations if become 1
                if(hap_iter->second->ancestor[j] == 0) {
                    if(hap_iter->first[j] == 0) {
                        assert(hap_iter->second->nnm[j] >= 1);
                        num_non_muts0 = hap_iter->second->nnm[j];
                        num_non_muts1 = hap_iter->second->nnm[j] - 1;
                    } else {
                        assert(hap_iter->first[j] == 1);
                        num_non_muts0 = hap_iter->second->nnm[j] + 1;
                        num_non_muts1 = hap_iter->second->nnm[j];
                    }
                } else {
                    assert(hap_iter->second->ancestor[j] == 1);
                    if(hap_iter->first[j] == 0) {
                        num_non_muts0 = hap_iter->second->nnm[j];
                        num_non_muts1 = hap_iter->second->nnm[j] + 1;
                    } else {
                        assert(hap_iter->first[j] == 1);
                        assert(hap_iter->second->nnm[j] >= 1);
                        num_non_muts0 = hap_iter->second->nnm[j] - 1;
                        num_non_muts1 = hap_iter->second->nnm[j];
                    }
                }

                int num_muts0 = hap_iter->second->num_at_table - num_non_muts0;
                int num_muts1 = hap_iter->second->num_at_table - num_non_muts1;

                double rightside0 = gsl_sf_gamma(alpha_h + (double)num_non_muts0) * gsl_sf_gamma(beta_h + (double)num_muts0) / gsl_sf_gamma(alpha_h + beta_h + (double)hap_iter->second->num_at_table);

                double rightside1 = gsl_sf_gamma(alpha_h + (double)num_non_muts1) * gsl_sf_gamma(beta_h + (double)num_muts1) / gsl_sf_gamma(alpha_h + beta_h + (double)hap_iter->second->num_at_table);

                double hap_goes_0 = leftside0 * rightside0; 
                double hap_goes_1 = leftside1 * rightside1; 
                double normalizing_constant = hap_goes_0 + hap_goes_1;
                double probability = hap_goes_0/normalizing_constant;

                double unif = gsl_rng_uniform(rg);

                int old_allele = hap_iter->first[j];
                if(unif < probability) {
                    // change hap_itj to 0
                    hap_iter->first[j] = 0;
                    u = tmp_u0;
                    up = tmp_up0;
                    udp = tmp_udp0;
                    // update ancestor's nnm
                    if(old_allele == 0) {
                        // no change
                        assert(hap_iter->second->nnm[j] == num_non_muts0);
                    } else {
                        assert(old_allele == 1);
                        if(hap_iter->second->ancestor[j] == 0) {
                            ++hap_iter->second->nnm[j];
                            assert(hap_iter->second->nnm[j] == num_non_muts0);
                        } else {
                            --hap_iter->second->nnm[j];
                            assert(hap_iter->second->nnm[j] == num_non_muts0);
                        }
                    }
                } else {
                    // change hap_itj to 1
                    hap_iter->first[j] = 1;
                    u = tmp_u1;
                    up = tmp_up1;
                    udp = tmp_udp1;
                    // update ancestor's nnm
                    if(old_allele == 1) {
                        // no change
                        assert(hap_iter->second->nnm[j] == num_non_muts1);
                    } else {
                        assert(old_allele == 0);
                        if(hap_iter->second->ancestor[j] == 1) {
                            ++hap_iter->second->nnm[j];
                            assert(hap_iter->second->nnm[j] == num_non_muts1);
                        } else {
                            --hap_iter->second->nnm[j];
                            assert(hap_iter->second->nnm[j] == num_non_muts1);
                        }
                    }
                }
            } // end iterating over sites
            ++hap_counter;
        }
        // Done updating parameters
        ///////////////////////////////////////////

        ///////////////////////////////////////////
        // Record summaries of posterior samples
        // compute posterior for each locus of the haplotypes individually as in Stevens et al
        for(uint i=0; i < num_haps; ++i) {
            for(uint j=0; j < segsites; ++j) {
                if(haplotypes[i].first[j] == 0) {
                    ++hap_counts0[i][j];
                } else {
                    ++hap_counts1[i][j];
                }
            }
        }

        // update posterior_count
        if(iteration >= burn_in) {
            for(uint ii=0; ii < num_haps/2; ++ii) {
                pair<vector<int>, vector<int> > cur_sample(haplotypes[2*ii].first, haplotypes[2*ii+1].first);
                if(posterior_counts[ii].find(cur_sample) == posterior_counts[ii].end()) {
                    posterior_counts[ii][cur_sample] = 1;
                } else {
                    ++posterior_counts[ii][cur_sample];
                }
            }
        }
    } // end Gibbs sampling
    ////////////////////////////////////////////////////

    // print out all ancestral haplotypes
    // infer haplotypes based on single locus approach
    for(uint i=0; i < num_haps; ++i) {
        for(uint j=0; j < segsites; ++j) {
            if(hap_counts0[i][j] > hap_counts1[i][j]) {
                inf_haps[i][j] = 0; 
            } else {
                inf_haps[i][j] = 1;
            }
        }
    }
    ///////////////////////////////////////////////////
    for(uint ii=0; ii < num_haps/2; ++ii) {
        int max_count = 0;
        pair<vector<int>, vector<int> > best_haps;
        for(boost::unordered_map<pair<vector<int>, vector<int> >, int>::iterator post_iter = posterior_counts[ii].begin(); post_iter != posterior_counts[ii].end(); ++post_iter) {
            if(post_iter->second > max_count) {
                best_haps = post_iter->first;
                max_count = post_iter->second;
            }
        }

        int firstwithfirst0 = 0;
        int num_heterozygous = 0;
        vector<int> best_hap0 = best_haps.second;
        vector<int> best_hap1 = best_haps.first;
        for(uint j=0; j<segsites; ++j) {
            if(genotypes[ii][j] == 2) {
                ++num_heterozygous;
                if(best_hap0[j] != data.pairings[ii].first[j]) {
                    ++firstwithfirst0;
                }
            }
        }
        int firstwithsecond0 = 0;
        for(uint j=0; j<segsites; ++j) {
            if(genotypes[ii][j] == 2) {
                if(best_hap0[j] != data.pairings[ii].second[j]) {
                    ++firstwithsecond0;
                }
            }
        }
          comparison_num_mistakes[ii] = min(firstwithfirst0, firstwithsecond0);
    }

    //int tot_num_mistakes = 0;
    //for(uint i=0; i < num_haps/2; ++i) {
    //    tot_num_mistakes += comparison_num_mistakes[i];
    //}

    int old_precision = cout.precision();
    cout << setprecision(16);

    // compute err_s
    int total_site_count = 0;
    int total_mistakes = 0;
    for(uint i=0; i < num_haps/2; ++i) {
        if(gen_het_num[i] > 1) {
            total_site_count += gen_het_num[i];
            total_mistakes += comparison_num_mistakes[i];
        }
    }
    double err_s = (double)total_mistakes/(double)total_site_count;

    cout << "err_s: " << err_s << endl;

    cout << "totalsitemistakes: " << total_mistakes << endl;
    cout << "totalsitecount: " << total_site_count << endl;

    int total_ind_count = 0;
    int total_ind_mistakes = 0;

    // compute err_i
    for(uint i=0; i < num_haps/2; ++i) {
        if(gen_het_num[i] > 1) {
            ++total_ind_count;
            if(comparison_num_mistakes[i] > 0) {
                ++total_ind_mistakes;
            }
        }
    }

    double err_i = (double)total_ind_mistakes/(double)total_ind_count;
    cout << "err_i: " << err_i << endl;
    cout << "totalindmistakes: " << total_ind_mistakes << endl;
    cout << "totalindcount: " << total_ind_count << endl;
    cout << setprecision(old_precision);
    
    return 0;
}
